'''
DeepLabV3Plus-Pytorch ONNX导出 - ResNet101 Cityscapes
----------------------------------------------
@作者: gx
@邮箱: gaoxukkk@qq.com
@创建日期：2024年11月14日
@修改日期：2024年11月14日 - 适配ResNet101 Cityscapes
'''

import torch
import torch.nn as nn
import cv2
import numpy as np
from PIL import Image
import os

from network.modeling import deeplabv3plus_resnet101

# BN层动量设置
def set_bn_momentum(model, momentum=0.1):
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.momentum = momentum

# Cityscapes颜色映射表
def cityscapes_cmap():
    """
    Cityscapes数据集的颜色映射表
    返回19个类别的RGB颜色值
    """
    # Cityscapes的19个类别颜色映射 (不包括ignore的类别)
    colors = [
        [128, 64, 128],   # road
        [244, 35, 232],   # sidewalk
        [70, 70, 70],     # building
        [102, 102, 156],  # wall
        [190, 153, 153],  # fence
        [153, 153, 153],  # pole
        [250, 170, 30],   # traffic light
        [220, 220, 0],    # traffic sign
        [107, 142, 35],   # vegetation
        [152, 251, 152],  # terrain
        [70, 130, 180],   # sky
        [220, 20, 60],    # person
        [255, 0, 0],      # rider
        [0, 0, 142],      # car
        [0, 0, 70],       # truck
        [0, 60, 100],     # bus
        [0, 80, 100],     # train
        [0, 0, 230],      # motorcycle
        [119, 11, 32],    # bicycle
    ]
    return np.array(colors, dtype=np.uint8)


class DeepLabV3PlusResNet101Cityscapes(torch.nn.Module):
    def __init__(self):
        super().__init__()

        # Cityscapes数据集
        self.num_classes = 19
        self.output_stride = 16

        self.model = deeplabv3plus_resnet101(num_classes=self.num_classes, output_stride=self.output_stride)
        # 修改模型中的BN层的动量
        set_bn_momentum(self.model.backbone, momentum=0.01)
        
        # 加载模型 - 请根据实际的checkpoint路径修改
        checkpoint_path = "checkpoints/best_deeplabv3plus_resnet101_cityscapes_os16.pth.tar"
        if os.path.exists(checkpoint_path):
            state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
            self.model.load_state_dict(state_dict["model_state"])
            print(f"Successfully loaded checkpoint from {checkpoint_path}")
        else:
            print(f"Warning: Checkpoint not found at {checkpoint_path}")
            print("Please download the Cityscapes ResNet101 checkpoint and place it in the checkpoints directory")
            print("Download link: https://drive.google.com/file/d/1t7TC8mxQaFECt4jutdq_NMnWxdm6B-Nb/view?usp=sharing")
    
    def forward(self, x):
        # 获取模型的原始输出 (logits)
        logits = self.model(x)
        return logits

        
device = "cpu"
model = DeepLabV3PlusResNet101Cityscapes().eval().to(device)


# ------- 数据导入 -------
# 使用Cityscapes的典型输入尺寸 1024x2048
image = cv2.imread("samples/CAM_FRONT.jpg")
if image is None:
    print("Warning: samples/CAM_FRONT.jpg not found, creating a dummy image for testing")
    # 创建一个测试图像
    image = np.random.randint(0, 255, (1024, 2048, 3), dtype=np.uint8)
else:
    image = cv2.resize(image, (2048, 1024))  # Cityscapes典型尺寸


# ------- 预处理 -------
# To RGB
image = image[..., ::-1]  # 是一种toRGB的方法 

#  Normalize the image (scale pixel values to [0, 1])
image = (image / 255.0).astype(np.float32)

# Normalize with given mean and std for each channel (RGB)
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
for c in range(image.shape[2]):  # Iterate over the RGB channels
    image[..., c] = (image[..., c] - mean[c]) / std[c]  # Standardize each channel
    
# Convert image to (C, H, W) format (channels first)
image = image.transpose(2, 0, 1)  # Shape: (H, W, C) -> (C, H, W)

# Add batch dimension (1, C, H, W)
image = image[None, ...]  # Add a new dimension for the batch (1, C, H, W)

# Convert to PyTorch tensor
image = torch.from_numpy(image).to(device)


# ------- 后处理 -------
cmap = cityscapes_cmap()
with torch.no_grad():
    logits = model(image)

    # 这里才可以进行 max 和 numpy 转换
    pred = logits.max(1)[1].cpu().numpy()[0]  # 获取预测的类别索引，并转为 NumPy 数组

    # 进行颜色映射
    colorized_preds = cmap[pred].astype('uint8')
    colorized_preds = Image.fromarray(colorized_preds)
    colorized_preds.save(os.path.join("samples/", 'resnet101_cityscapes_res.png'))


# ------- 导出 ONNX 模型 -------
import onnx

input_layer_names = ["images"]
output_layer_names = ["output"]
model_path = "checkpoints/deeplabv3plus_resnet101_cityscapes.onnx"

# 导出模型
print(f'Starting export with onnx {onnx.__version__}.')
torch.onnx.export(
    model,  # 模型对象
    (image,),  # 模型的输入（输入数据）
    f=model_path,  # 保存的 ONNX 文件路径
    verbose=False,  # 是否打印详细信息
    opset_version=12,  # 使用 opset 版本 12
    do_constant_folding=True,  # 是否进行常量折叠优化
    input_names=input_layer_names,  # 输入层的名字
    output_names=output_layer_names,  # 输出层的名字
    dynamic_axes=None  # 是否有动态维度
)

# ------- 检查 ONNX 模型 -------
model_onnx = onnx.load(model_path)  # 加载 ONNX 模型
onnx.checker.check_model(model_onnx)  # 检查 ONNX 模型

# ------- 简化 ONNX -------
import onnxsim
print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.')
model_onnx, check = onnxsim.simplify(
    model_onnx,
    dynamic_input_shape=False,
    input_shapes=None)
assert check, 'assert check failed'
onnx.save(model_onnx, model_path)

print(f'Onnx model saved as {model_path}')


print("ResNet101 Cityscapes ONNX export completed!") 